package main

import (
	"context"
	"log"
	"mag/internal/database"
	"mag/internal/utils"
	"math"
	"time"

	"github.com/emirpasic/gods/sets/hashset"
	"go.mongodb.org/mongo-driver/bson"
	"go.mongodb.org/mongo-driver/mongo"
	"go.mongodb.org/mongo-driver/mongo/options"
)

func main() {
	// 初始化容器
	// originInDict := make(map[int64][]int64)
	// originOutDict := make(map[int64][]int64)
	log.Println("开始导入")
	rangInMap_1 := make(map[int64]*hashset.Set)
	rangOutMap_1 := make(map[int64]*hashset.Set)

	rangInMap_2 := make(map[int64]*hashset.Set)
	rangOutMap_2 := make(map[int64]*hashset.Set)

	rangInMap_3 := make(map[int64]*hashset.Set)
	rangOutMap_3 := make(map[int64]*hashset.Set)

	ctx := context.Background()
	client := database.InitMongo()
	collection := client.Database("mag2020").Collection("pageinfo")

	cur, err := collection.Find(ctx, bson.M{})
	if err != nil {
		log.Fatal(err)
	}
	defer cur.Close(ctx)

	for cur.Next(ctx) {

		var result database.PageInfo
		err := cur.Decode(&result)
		if err != nil {
			log.Fatal(err)
		}
		// 开始过滤一遍
		if len(result.Out) == 0 && len(result.In) == 0 {
			continue
		}
		// 初始化 out data
		rangOutMap_1[result.Id] = hashset.New()
		rangOutMap_2[result.Id] = hashset.New()
		rangOutMap_3[result.Id] = hashset.New()

		if result.Year < 1950 {
			for _, item := range result.Out {
				rangOutMap_1[result.Id].Add(item)
				rangOutMap_2[result.Id].Add(item)
				rangOutMap_3[result.Id].Add(item)

				if rangInMap_1[item] == nil {
					rangInMap_1[item] = hashset.New()
				}
				if rangInMap_2[item] == nil {
					rangInMap_2[item] = hashset.New()
				}
				if rangInMap_3[item] == nil {
					rangInMap_3[item] = hashset.New()
				}
				rangInMap_1[item].Add(result.Id)
				rangInMap_2[item].Add(result.Id)
				rangInMap_3[item].Add(result.Id)
			}
		} else {

			flag := math.Pow(0.9892, float64(result.Year-1949))
			for _, item := range result.Out {
				// 随机3次概率
				if utils.IsInRange(flag) {
					rangOutMap_1[result.Id].Add(item)
					if rangInMap_1[item] == nil {
						rangInMap_1[item] = hashset.New()
					}
					rangInMap_1[item].Add(result.Id)
				}
				if utils.IsInRange(flag) {
					rangOutMap_2[result.Id].Add(item)
					if rangInMap_2[item] == nil {
						rangInMap_2[item] = hashset.New()
					}
					rangInMap_2[item].Add(result.Id)
				}
				if utils.IsInRange(flag) {
					rangOutMap_3[result.Id].Add(item)
					if rangInMap_3[item] == nil {
						rangInMap_3[item] = hashset.New()
					}
					rangInMap_3[item].Add(result.Id)
				}
			}
		}
	}

	log.Println("数据导入完毕，开始计算 D")
	// 开始计算 disruption
	outchan := calDisruption(rangOutMap_1, rangInMap_1)
	go handleOutChan(outchan, "disruption_random_1")

	outchan1 := calDisruption(rangOutMap_2, rangInMap_2)
	go handleOutChan(outchan1, "disruption_random_2")

	outchan2 := calDisruption(rangOutMap_3, rangInMap_3)
	handleOutChan(outchan2, "disruption_random_3")

	log.Println("over")
	time.Sleep(99999)
}

type disruptionItem struct {
	ID int64   `bson:"_id"`
	D  float64 `bson:"d"`
	I  int64   `bson:"i"`
	J  int64   `bson:"j"`
	K  int64   `bson:"k"`
}

func (c *disruptionItem) CalD() {
	c.D = (float64(c.I) - float64(c.J)) / (float64(c.I) + float64(c.J) + float64(c.K))
}

func calDisruption(outMap, inMap map[int64]*hashset.Set) chan *disruptionItem {
	outChan := make(chan *disruptionItem, 10000)
	go func() {
		for key, outSet := range outMap {
			retData := disruptionItem{ID: key}

			// 获取 linksin set, 如果没有了，就跳过
			inSet := inMap[key]
			if inSet == nil || outSet.Size() == 0 {
				continue
			}

			outInSet := hashset.New()

			for _, item := range outSet.Values() {
				if outinset, ok := inMap[item.(int64)]; ok {
					outInSet.Add(outinset.Values()...)
				}
			}

			// 计算颠覆度参数
			for _, nodeID := range inSet.Values() {
				nodeIDOutSet := outMap[nodeID.(int64)]
				if Is_intersection(nodeIDOutSet, outSet) {
					retData.J += 1
				} else {
					retData.I += 1
				}
			}
			// 排除自身节点
			outInSet.Remove(key)
			retData.K = Cal_uniq_count(outInSet, inSet)
			retData.CalD()
			outChan <- &retData
		}
		close(outChan)
	}()
	return outChan
}

func handleOutChan(outchan chan *disruptionItem, collectionName string) {

	collection := database.InitMongo().Database("mag2020").Collection(collectionName)

	opts := options.BulkWrite().SetOrdered(false)
	models := []mongo.WriteModel{}
	i := 0
	for item := range outchan {
		i += 1
		models = append(models, mongo.NewInsertOneModel().SetDocument(item))
		if i%10000 == 0 {
			retTmp, err := collection.BulkWrite(context.Background(), models, opts)
			if err != nil {
				log.Println("bulk upsert fail", err)
				log.Printf("%+v", retTmp)
			}
			models = []mongo.WriteModel{}
		}
	}
	if len(models) > 0 {
		retTmp, err := collection.BulkWrite(context.Background(), models, opts)
		if err != nil {
			log.Println("bulk upsert fail", err)
			log.Printf("%+v", retTmp)
		}
	}
}

// 计算是否有交集
func Is_intersection(a, b *hashset.Set) bool {
	for _, v := range a.Values() {
		if b.Contains(v.(int64)) {
			return true
		}
	}
	return false
}

// 计算差
func Cal_uniq_count(a, b *hashset.Set) (count int64) {
	for _, v := range a.Values() {
		if !b.Contains(v.(int64)) {
			count += 1
		}
	}
	return
}
